//
// Created by 王有东 on 2024/4/7.
//
//////////////////////////////
#include <iostream>
#include <cassert>
#include <vector>
#include <stack>
using namespace std;
///红色
#define RED 1
///黑色
#define BLACK 2

///////////////////////////////
//定义红黑树的节点应该包含:
// 1. 树存的值_val
// 2. 树的颜色 _color
// 3. 指向孩子的指针 _left, _right
template <class T>
//必须定义为public，成员可以直接被使用
struct RBTreeNode
{
    RBTreeNode(const T& val)
        : _val(val)
        , _color(RED)
        , _left(nullptr)
        , _right(nullptr)
        , _parent(nullptr)
    {}

    typedef RBTreeNode<T> Node;
    T _val;
    int _color;
    Node* _left;
    Node* _right;
    Node* _parent;
};

////////////////////////////////
//定义红黑树，成员变量必须是节点的指针，以便创造节点
//且为了可维护性，要将成员变量定义为private
template <class T>
class RBTree
{
public:
    typedef RBTreeNode<T> Node;

    void InOrder()
    {
        _InOrder(_root);
        cout << endl;
    }

    //插入节点，按照二叉搜索树来
    bool insert(const T& val)
    {
        //空树
        if(_root == nullptr)
        {
            _root = new Node(val);
            _root->_color = BLACK;
            return true;
        }
        else
        {
            Node* cur = _root;
            Node* parent = cur;
            while(cur)
            {
                if(cur->_val < val)
                {
                    parent = cur;
                    cur = cur->_right;
                }
                else if(cur->_val > val)
                {
                    parent = cur;
                    cur = cur->_left;
                }
                else
                    return false;
            }

            //找到可以插入的位置cur
            cur = new Node(val);
            if(parent->_val < val)
            {
                parent->_right = cur;
            }
            else
            {
                parent->_left = cur;
            }
            cur->_parent = parent;

            //判断颜色,并改正
            AssertColor(parent, cur);

        }
        return true;
    }
    bool Check()
    {
        ////////非递归版本
        //return CheckRoot() && CheckColor(_root) && CheckNumB(_root);

        ////////递归版本
        Node* cur = _root;
        int num = 0;
        while(cur)
        {
            if(cur->_color == BLACK)
                ++num;
            cur = cur->_left;
        }

        return CheckRoot() && CheckColor(_root)
            && CheckNumB_R(_root, num, 0);
    }

private:
    Node* _root = nullptr;
    //递归版本
    bool CheckNumB_R(Node* root, int num, int count)
    {
        if(root == nullptr)
        {
            if(num != count)
                return false;
            return true;
        }

        if(root->_color == BLACK)
            ++count;

        return CheckNumB_R(root->_left, num, count) && CheckNumB_R(root->_right, num, count);
    }
    //非递归版本
    bool CheckNumB(Node* root)
    {
        vector<Node*> check; //检查是否右子树已经遍历过了，要记录所有走过的右侧的节点
        stack<Node*> s;      //存路径
        Node* cur = root;
        Node* pre = nullptr;
        int num = 0;         //记录黑色节点的个数
        int copy = -1;       //记录前一条路径的黑色节点的个数

        while(cur || !s.empty())
        {
            int flag = 0;

            while(cur)  //遍历每条路径
            {
                if(cur->_color == BLACK)
                {
                    ++num;
                }
                s.push(cur);
                cur = cur->_left;
                flag = 1;
            }

            if(flag == 1) //说明走的路径是新的
            {
                if(copy != -1)
                {
                    if(copy != num) //前一个路径的黑色节点个数不等于当前路径的黑色节点个数，不满足条件，返回false
                        return false;
                }
                copy = num;
            }

            if(s.top()->_right == nullptr || CheckRepeated(check, s.top()->_right))  //右子树为空 / 右子树走过了，路径结束，换路径
            {
                if(s.top()->_color == BLACK)
                {
                    --num;
                }
                s.pop();
            }
            else //右子树不为空，且没有走过，就去遍历右子树
            {
                cur = s.top()->_right;
                pre = cur;
                check.push_back(pre);
            }
        }
        //遍历完每条路径后，黑色节点个数都相等
        return true;

    }
    //检查右子树是否已经遍历过了
    bool CheckRepeated(vector<Node*>& check, Node*& pre)
    {
        for(auto e : check)
        {
            if(pre == e)
            {
                return true;
            }
        }
        return false;
    }


    //检查根是否为黑色
    bool CheckRoot()
    {
        if(_root == nullptr)
            return true;

        return _root->_color == BLACK;
    }

    //检查是否有连续的红色节点
    bool CheckColor(Node* root)
    {
        if(root == nullptr)
            return true;
        return CheckColor(root->_left)
        && CheckColor(root->_right)
        && (root->_color == BLACK || root->_parent->_color == BLACK);
    }

    void AssertColor(Node*& parent, Node*& cur)
    {

        while(parent && parent != _root)
        {
            Node* grandfather = parent->_parent;
            Node* uncle = (parent == grandfather->_left) ? grandfather->_right : grandfather->_left;
           //插入位置的父亲为红色节点，需要处理
           if(cur->_color == RED && parent->_color == RED)
           {
               //叔叔存在，且为红
               if(uncle && uncle->_color == RED)
               {
                   parent->_color = uncle->_color = BLACK;
                   grandfather->_color = RED;

                   cur = grandfather;
                   parent = cur->_parent;
               }
               //叔叔存在,且为黑 / 叔叔不存在  旋转+变色
               else
               {
                   if(uncle == grandfather->_left)
                   {
                       if(cur == parent->_right)
                       {
                           RotationL(grandfather);
                           parent->_color = BLACK;
                           grandfather->_color = RED;
                       }
                       else
                       {
                           RotationR(parent);
                           RotationL(grandfather);
                           cur->_color = BLACK;
                           grandfather->_color = RED;
                       }
                   }
                   else
                   {
                       if(cur == parent->_left)
                       {
                           RotationR(grandfather);
                           parent->_color = BLACK;
                           grandfather->_color = RED;
                       }
                       else
                       {
                           RotationL(parent);
                           RotationR(grandfather);
                           cur->_color = BLACK;
                           grandfather->_color = RED;
                       }
                   }
                   break;
               }
           }
           //插入位置的父亲为黑色节点，不需要处理
           else
           {
               break;
           }
       }

        _root->_color = BLACK;
    }
    void RotationL(Node* parent)
    {
        Node* pR = parent->_right;
        Node* pRL = pR->_left;

        parent->_right = pRL;
        if(pRL)
            pRL->_parent = parent;
        pR->_left = parent;
        Node* ppnode = parent->_parent;
        parent->_parent = pR;

        if(parent == _root)
        {
            _root = pR;
            pR->_parent = nullptr;
        }
        else
        {
            if(ppnode->_left == parent)
                ppnode->_left = pR;
            else
                ppnode->_right = pR;
            pR->_parent = ppnode;
        }

    }
    void RotationR(Node* parent)
    {
        Node* pL = parent->_left;
        Node* pLR = pL->_right;

        parent->_left = pLR;
        if(pLR)
            pLR->_parent = parent;
        pL->_right = parent;
        Node* ppnode = parent->_parent;
        parent->_parent = pL;

        if(_root == parent)
        {
            _root = pL;
            pL->_parent = nullptr;
        }
        else
        {
            if(ppnode->_left == parent)
                ppnode->_left = pL;
            else
                ppnode->_right = pL;
            pL->_parent = ppnode;
        }

    }
    void _InOrder(Node* root)
    {
        if(root == nullptr)
            return;
        _InOrder(root->_left);
        cout << root->_val << " ";
        _InOrder(root->_right);
    }
};






























